package rbTree;
import adelsonVelskyLandis.AVL;

import static rbTree.COLOR.BLACK;
import static rbTree.COLOR.RED;

public class RBTree {
    private static class RBTreeNode {
        private RBTreeNode parent;
        private RBTreeNode left;
        private RBTreeNode right;
        private COLOR color;
        int val ;

        public RBTreeNode(int val) {
            this.val = val;
            // 新节点默认都为 红色
            this.color = RED;
        }
    }

    public RBTreeNode root ;

    // 插入一个新节点
    public boolean insert(int val) {
        RBTreeNode newNode = new RBTreeNode(val);
        if (root == null) {
            root = newNode;
            root.color = BLACK; // 根节点颜色为 黑
            return true;
        }

        // 寻找插入位置
        RBTreeNode parent = null;
        RBTreeNode cur = root;
        while (cur != null) {
            if (cur.val < val) {
                parent = cur;
                cur = cur.right;
            } else if (cur.val > val) {
                parent = cur;
                cur = cur.left;
            } else {
                return true; // 重复节点
            }
        }
        // 至此 cur 为空，parent 为 cur 的父节点
        // 完成节点的插入
        if (parent.val > val) {
            parent.left = newNode;
        } else {
            parent.right = newNode;
        }
        newNode.parent = parent; // 双向连接
        // parent = cur.parent;
        cur = newNode;

        // 至此开始「向上调整颜色」
        // 新插入的结点是红色的, 如果父亲结点还是红色的, 就需要调整颜色
        while (parent != null && parent.color == RED) {
            // 定义祖父结点
            RBTreeNode grandparent = parent.parent; // 因为根节点必须是黑色, 所以祖父结点不可能为空
            // 第一种情况, parent 是 grandparent 的 左孩子
            if (parent == grandparent.left) {
                RBTreeNode uncle = grandparent.right;
                // 情况 1.1: parent 红色, uncle 不为空, 并且红色
                if (uncle != null && uncle.color == RED) {
                    parent.color = BLACK;
                    uncle.color = BLACK;
                    grandparent.color = RED; // 修改颜色

                    cur = grandparent;
                    parent = cur.parent;
                } else {
                    // 这里有还有两种情况
                    // 情况 1.3: cur 是 parent 的右孩子, 并且 uncle 为空, 或者 uncle 为黑色
                    if (cur == parent.right) {
                        // 这时候需要先左旋
                        rotateLeft(parent);
                        RBTreeNode temp = parent;
                        parent = cur;
                        cur = temp;
                    }

                    // 情况 1.2: cur 是 parent 的左孩子, 并且 uncle 为空, 或者 uncle 为黑色
                    // 需要 向右旋转
                    rotateRight(grandparent);
                    // 再修改颜色
                    grandparent.color = RED;
                    parent.color = BLACK;
                }
            } else {
                // parent = grandparent.right;
                RBTreeNode uncle = grandparent.left;
                if (uncle != null && uncle.color == RED) {
                    parent.color = BLACK;
                    uncle.color = BLACK;
                    grandparent.color = RED; // 先变为 红色, 方便后续操作, 最后再统一变成 黑色
                } else {
                    if (cur == parent.left) {
                        rotateRight(parent);
                        RBTreeNode temp = parent;
                        parent = cur;
                        cur = temp;
                    }

                    rotateLeft(grandparent);
                    grandparent.color = RED;
                    parent.color = BLACK;
                }
            }
        }

        // 最后再将根节点统一修改为 黑色
        root.color = BLACK;
        return true;
    }

    private void rotateLeft(RBTreeNode parent) {
        RBTreeNode rson = parent.right;
        RBTreeNode rsonLeft = rson.left;
        RBTreeNode grandparent = parent.parent;

        parent.right = rsonLeft;
        if (rsonLeft != null) {
            rsonLeft.parent = parent;
        }
        parent.parent = rson;
        rson.left = parent;

        if (root == parent) {
            root = rson;
            rson.parent = null;
        } else {
            if (grandparent.right == parent) {
                grandparent.right = rson;
            } else {
                grandparent.left = rson;
            }
            rson.parent = grandparent;
        }
    }

    private void rotateRight(RBTreeNode parent) {
        RBTreeNode lson = parent.left;
        RBTreeNode lsonRight = lson.right;
        RBTreeNode grandparent = parent.parent;

        parent.left = lsonRight;
        if (lsonRight != null) {
            lsonRight.parent = parent;
        }
        lson.right = parent;
        parent.parent = lson;

        if (root == parent) {
            root = lson;
            lson.parent = null;
        } else {
            if (grandparent.left == parent) {
                grandparent.left = lson;
            } else {
                grandparent.right = lson;
            }
            lson.parent = grandparent;
        }
    }

    // 判断一棵树 是不是 红黑树
    public boolean isRBTree() {
        if (root == null)   return true;

        if (root.color != BLACK) {  // 根节点必须为 黑色
            return false;
        }

        return checkRedNode(root) && checkBlackNode(root, 0, -1);
    }

    // 判断有没有 两个连续 的红色节点
    public boolean checkRedNode(RBTreeNode root) {
        if (root == null)   return true;
        if (root.color == RED) {
            if (root.parent != null && root.parent.color == RED) {
                return false;
            }
        }
        return checkRedNode(root.left) && checkRedNode(root.right);
    }


    /**
     * 检查黑色的结点符不符合要求
     * @param root 结点
     * @param pathBlackNum 当前路径的 黑色结点数, 刚开始默认是 -1, 第一次到达更结点的时候更新 neededBlackNum
     * @param neededBlackNum 总共需要的黑色结点数量
     */
    public boolean checkBlackNode(RBTreeNode root, int pathBlackNum, int neededBlackNum) {
        if (root == null)   return true;

        if (root.color == BLACK) {
            pathBlackNum ++;
        }
        // 到达根节点的时候检查一下
        if (root.left == null && root.right == null) {
            if (neededBlackNum == -1) { // 第一次走完一条完整路径的时候, 更新一下 neededBlackNum
                neededBlackNum = pathBlackNum;
            } else { // 更新了
                if (neededBlackNum != pathBlackNum) {
                    return false;
                }
            }
        }

        return checkBlackNode(root.left, pathBlackNum, neededBlackNum) && checkBlackNode(root.right, pathBlackNum, neededBlackNum);
    }

    public void inorder(RBTreeNode root) {
        if(root == null) {
            return;
        }
        inorder(root.left);
        System.out.print(root.val+" ");
        inorder(root.right);
    }
}
